package com.example.demo1.samplews;

import com.alibaba.fastjson.JSONObject;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;

import java.net.URI;
import java.util.ArrayList;
import java.util.List;

public class AiAgentWs extends WebSocketClient {

    String model;

    int maxRound;

    String firstSentence;

    String promptSystem;

    List<JSONObject> messages = new ArrayList<>();

    AiAgentWs partner;

    public AiAgentWs(String url, String model, int maxRound, String firstSentence, String promptSystem) {
        super(URI.create(url));
        this.model = model;
        if (maxRound < 2) {
            throw new RuntimeException("对话轮数不能小于3");
        }
        this.maxRound = maxRound;
        this.firstSentence = firstSentence;
        this.promptSystem = promptSystem;
        this.connect();
        this.setConnectionLostTimeout(10000);
    }

    public void start() {
        if (this.getPartner() == null) {
            throw new RuntimeException("没有可对话的aiAgent");
        }
        this.onMessage(firstSentence);
    }

    @Override
    public void onOpen(ServerHandshake serverHandshake) {
        System.out.println("【" + model + "】 已激活");
    }

    @Override
    public void onMessage(String s) {
        System.out.println("【" + model + "】： " + s);
        JSONObject userAgentSaid = new JSONObject();
        userAgentSaid.put("role","user");
        userAgentSaid.put("content", s);
        if (this.messages.size() > maxRound) {
            this.messages.remove(this.messages.get(1));
        }
        this.messages.add(userAgentSaid);
        JSONObject param = new JSONObject();
        param.put("model", partner.getModel().trim());
        param.put("messages", this.messages);
        param.put("stream", false);
        partner.send(param.toJSONString());
    }

    @Override
    public void send(String text) {
        JSONObject assistantAgentSaid = new JSONObject();
        assistantAgentSaid.put("role", "assistant");
        assistantAgentSaid.put("content", assistantAgentSaid);
        if (this.messages.size() > maxRound) {
            this.messages.remove(this.messages.get(1));
        }
        this.messages.add(assistantAgentSaid);
        super.send(text);
    }

    @Override
    public void onClose(int i, String s, boolean b) {
        System.out.println("【" + model + "】 已退出");
    }

    @Override
    public void onError(Exception e) {
        e.printStackTrace();
        System.out.println("【" + model + "】 出现故障");
    }

    public String getModel() {
        return model;
    }

    public void setModel(String model) {
        this.model = model;
    }

    public int getMaxRound() {
        return maxRound;
    }

    public void setMaxRound(int maxRound) {
        this.maxRound = maxRound;
    }

    public String getFirstSentence() {
        return firstSentence;
    }

    public void setFirstSentence(String firstSentence) {
        this.firstSentence = firstSentence;
    }

    public String getPromptSystem() {
        return promptSystem;
    }

    public void setPromptSystem(String promptSystem) {
        this.promptSystem = promptSystem;
    }

    public List<JSONObject> getMessages() {
        return messages;
    }

    public void setMessages(List<JSONObject> messages) {
        this.messages = messages;
    }

    public AiAgentWs getPartner() {
        return partner;
    }

    public void setPartner(AiAgentWs partner) {
        JSONObject promptSystemObj = new JSONObject();
        promptSystemObj.put("role", "system");
        promptSystemObj.put("content", partner.getPromptSystem());
        this.messages.add(promptSystemObj);
        this.partner = partner;
    }

}
